/*
 * Copyright © 2022-now 尛飛俠（Denvie） All rights reserved.
 */

package cn.denvie.sentinel.persistent.pull.file;

import cn.denvie.sentinel.persistent.pull.file.assist.RuleFileUtils;
import cn.denvie.sentinel.persistent.pull.file.assist.RuleListConverterUtils;
import cn.denvie.sentinel.persistent.pull.file.ext.DefaultStoreHandler;
import cn.denvie.sentinel.persistent.pull.file.ext.StoreHandler;
import com.alibaba.csp.sentinel.adapter.gateway.common.api.ApiDefinition;
import com.alibaba.csp.sentinel.adapter.gateway.common.api.GatewayApiDefinitionManager;
import com.alibaba.csp.sentinel.adapter.gateway.common.command.UpdateGatewayApiDefinitionGroupCommandHandler;
import com.alibaba.csp.sentinel.adapter.gateway.common.command.UpdateGatewayRuleCommandHandler;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayFlowRule;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayRuleManager;
import com.alibaba.csp.sentinel.command.handler.ModifyParamFlowRulesCommandHandler;
import com.alibaba.csp.sentinel.datasource.FileRefreshableDataSource;
import com.alibaba.csp.sentinel.datasource.FileWritableDataSource;
import com.alibaba.csp.sentinel.datasource.ReadableDataSource;
import com.alibaba.csp.sentinel.datasource.WritableDataSource;
import com.alibaba.csp.sentinel.init.InitFunc;
import com.alibaba.csp.sentinel.init.InitOrder;
import com.alibaba.csp.sentinel.slots.block.authority.AuthorityRule;
import com.alibaba.csp.sentinel.slots.block.authority.AuthorityRuleManager;
import com.alibaba.csp.sentinel.slots.block.degrade.DegradeRule;
import com.alibaba.csp.sentinel.slots.block.degrade.DegradeRuleManager;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.param.ParamFlowRuleManager;
import com.alibaba.csp.sentinel.slots.system.SystemRule;
import com.alibaba.csp.sentinel.slots.system.SystemRuleManager;
import com.alibaba.csp.sentinel.transport.util.WritableDataSourceRegistry;

import java.io.FileNotFoundException;
import java.util.Iterator;
import java.util.List;
import java.util.ServiceLoader;
import java.util.Set;
import java.util.logging.Logger;

/**
 * InitFunc 实现类，处理文件数据源初始化逻辑。
 *
 * @author Denvie
 * @date 2022/4/5
 * @since 1.0.0
 */
@InitOrder(-1)
public class FileDataSourceInitFunc implements InitFunc {
    private static final Logger logger = Logger.getLogger(FileDataSourceInitFunc.class.getName());

    private StoreHandler storeHandler;

    @Override
    public void init() throws Exception {
        logger.info("FileDataSourceInitFunc.init() start");

        // 创建规则文件存储目录
        this.storeHandler = loadStoreHandler();
        String rulePath = storeHandler.storeDir();
        RuleFileUtils.mkdirIfNotExits(rulePath);
        logger.info("Sentinel rules store dir: " + rulePath);

        // 创建规则文件
        RuleFileUtils.createFileIfNotExits(this.storeHandler.rulesMap());

        // 处理流控规则
        handleFlowRules();

        // 处理降级规则
        handleDegradeRules();

        // 处理系统规则
        handleSystemRules();

        // 处理热点参数规则
        handleParamFlowRules();

        // 处理授权规则
        handleAuthRules();

        // 处理网关规则
        handleGatewayRules();

        // 处理网关Api分组规则
        handleApiDefinitionRules();
    }

    private StoreHandler loadStoreHandler() {
        ServiceLoader<StoreHandler> storeHandlers = ServiceLoader.load(StoreHandler.class);
        Iterator<StoreHandler> iterator = storeHandlers.iterator();
        if (iterator.hasNext()) {
            return iterator.next();
        }
        return new DefaultStoreHandler();
    }

    private void handleFlowRules() throws FileNotFoundException {
        String ruleFilePath = this.storeHandler.rulesMap().get(StoreHandler.FLOW_RULE_PATH);

        // 创建流控规则的可读数据源
        // FileRefreshableDataSource 会周期性的读取文件以获取规则
        ReadableDataSource<String, List<FlowRule>> readableDataSource = new FileRefreshableDataSource<>(
                ruleFilePath, RuleListConverterUtils.flowRuleListParser);

        // 将可读数据源注册至 FlowRuleManager，这样当规则文件发生变化时，就会更新规则到内存
        FlowRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<List<FlowRule>> writableDataSource = new FileWritableDataSource<>(
                ruleFilePath, RuleListConverterUtils.flowRuleEncoder);

        // 将可写数据源注册至 transport 模块的 WritableDataSourceRegistry 中.
        // 这样收到控制台推送的规则时，Sentinel 会先更新到内存，然后将规则写入到文件中.
        WritableDataSourceRegistry.registerFlowDataSource(writableDataSource);
    }

    private void handleDegradeRules() throws FileNotFoundException {
        String degradeRuleFilePath = this.storeHandler.rulesMap().get(StoreHandler.DEGRADE_RULE_PATH);
        ReadableDataSource<String, List<DegradeRule>> readableDataSource = new FileRefreshableDataSource<>(
                degradeRuleFilePath, RuleListConverterUtils.degradeRuleListParser);
        DegradeRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<List<DegradeRule>> writableDataSource = new FileWritableDataSource<>(
                degradeRuleFilePath, RuleListConverterUtils.degradeRuleEncoder);
        WritableDataSourceRegistry.registerDegradeDataSource(writableDataSource);
    }

    private void handleSystemRules() throws FileNotFoundException {
        String systemRuleFilePath = this.storeHandler.rulesMap().get(StoreHandler.SYSTEM_RULE_PATH);
        ReadableDataSource<String, List<SystemRule>> readableDataSource = new FileRefreshableDataSource<>(
                systemRuleFilePath, RuleListConverterUtils.sysRuleListParser);
        SystemRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<List<SystemRule>> writableDataSource = new FileWritableDataSource<>(
                systemRuleFilePath, RuleListConverterUtils.sysRuleEncoder);
        WritableDataSourceRegistry.registerSystemDataSource(writableDataSource);
    }


    private void handleParamFlowRules() throws FileNotFoundException {
        String paramFlowRuleFilePath = this.storeHandler.rulesMap().get(StoreHandler.HOT_PARAM_RULE);
        ReadableDataSource<String, List<ParamFlowRule>> readableDataSource = new FileRefreshableDataSource<>(
                paramFlowRuleFilePath, RuleListConverterUtils.paramFlowRuleListParser);
        ParamFlowRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<List<ParamFlowRule>> writableDataSource = new FileWritableDataSource<>(
                paramFlowRuleFilePath, RuleListConverterUtils.paramRuleEncoder);
        ModifyParamFlowRulesCommandHandler.setWritableDataSource(writableDataSource);
    }

    private void handleAuthRules() throws FileNotFoundException {
        String authFilePath = this.storeHandler.rulesMap().get(StoreHandler.AUTH_RULE_PATH);
        ReadableDataSource<String, List<AuthorityRule>> readableDataSource = new FileRefreshableDataSource<>(
                authFilePath, RuleListConverterUtils.authorityRuleParser);
        AuthorityRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<List<AuthorityRule>> writableDataSource = new FileWritableDataSource<>(
                authFilePath, RuleListConverterUtils.authorityEncoder);
        WritableDataSourceRegistry.registerAuthorityDataSource(writableDataSource);
    }

    private void handleGatewayRules() throws FileNotFoundException {
        String gatewayFilePath = this.storeHandler.rulesMap().get(StoreHandler.GATEWAY_RULE);
        ReadableDataSource<String, Set<GatewayFlowRule>> readableDataSource = new FileRefreshableDataSource<>(
                gatewayFilePath, RuleListConverterUtils.gatewayParser);
        GatewayRuleManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<Set<GatewayFlowRule>> writableDataSource = new FileWritableDataSource<>(
                gatewayFilePath, RuleListConverterUtils.gatewayEncoder);
        UpdateGatewayRuleCommandHandler.setWritableDataSource(writableDataSource);
    }

    private void handleApiDefinitionRules() throws FileNotFoundException {
        String apiFilePath = this.storeHandler.rulesMap().get(StoreHandler.API_DEFINITION_RULE);
        ReadableDataSource<String, Set<ApiDefinition>> readableDataSource = new FileRefreshableDataSource<>(
                apiFilePath, RuleListConverterUtils.apiDefinitionParser);
        GatewayApiDefinitionManager.register2Property(readableDataSource.getProperty());

        WritableDataSource<Set<ApiDefinition>> writableDataSource = new FileWritableDataSource<>(
                apiFilePath, RuleListConverterUtils.apiDefinitionEncoder);
        UpdateGatewayApiDefinitionGroupCommandHandler.setWritableDataSource(writableDataSource);
    }
}
